import streamlit as st
from tkinter import Tk, filedialog
import os

# 初始化 session state
if "folder" not in st.session_state:
    st.session_state["folder"] = None
if "selected_blocks" not in st.session_state:
    st.session_state["selected_blocks"] = set()

# --- 工具函数 ---

def select_directory():
    root = Tk()
    root.withdraw()
    return filedialog.askdirectory()

def extract_blocks(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    blocks = []
    i = 0
    while i < len(lines):
        line = lines[i]
        if line.strip().startswith("@ApiModelProperty"):
            api_line = line
            if i + 1 < len(lines) and lines[i + 1].strip().startswith("private"):
                private_line = lines[i + 1]
                blocks.append(api_line + private_line)
                i += 2
            else:
                i += 1
        else:
            i += 1
    return blocks

# --- 页面布局 ---

st.set_page_config(page_title="Java 字段提取工具", layout="wide")
st.title("📦 Java @ApiModelProperty 字段提取工具")

# 选择文件夹
if st.button("📁 选择文件夹"):
    selected_folder = select_directory()
    st.session_state["folder"] = selected_folder
    st.session_state["selected_blocks"] = set()  # 清空旧选择

# 加载和展示文件
if st.session_state["folder"]:
    folder = st.session_state["folder"]
    st.success(f"当前目录：{folder}")

    java_files = [f for f in os.listdir(folder) if f.endswith(".java")]
    file_paths = {f: os.path.join(folder, f) for f in java_files}

    selected_files = st.multiselect(
        "选择要处理的 Java 文件：", java_files, default=java_files
    )

    all_blocks = {}
    for file in selected_files:
        path = file_paths[file]
        blocks = extract_blocks(path)
        all_blocks[file] = blocks

    # 文件中提取字段展示并选择
    st.markdown("### 📝 文件字段选择")
    for file, blocks in all_blocks.items():
        if blocks:
            st.markdown(f"**📄 {file}**")
            block_labels = [f"{file} | {block.strip().replace(chr(10), ' ')}" for block in blocks]
            selected = st.multiselect(
                f"从 {file} 中选择字段", block_labels, key=f"select_{file}"
            )
            for s in selected:
                st.session_state["selected_blocks"].add(s)
            # 如果用户取消了某些已选项，也从 set 中移除
            for s in block_labels:
                if s not in selected and s in st.session_state["selected_blocks"]:
                    st.session_state["selected_blocks"].discard(s)

    # 搜索功能：支持多个关键词（空格或逗号分隔）
    st.markdown("### 🔍 搜索字段（支持多个关键词）")
    raw_keyword = st.text_input("输入关键字（空格或英文逗号分隔）：", "")

    search_matches = []
    keywords = [k.strip().lower() for k in raw_keyword.replace(',', ' ').split() if k.strip()]

    if keywords:
        for file, blocks in all_blocks.items():
            for block in blocks:
                if any(k in block.lower() for k in keywords):
                    label = f"{file} | {block.strip().replace(chr(10), ' ')}"
                    search_matches.append(label)

    # 默认选中的是 session 中已有的
    default_selected = [s for s in search_matches if s in st.session_state["selected_blocks"]]

    selected_search = st.multiselect(
        "搜索结果选择：", options=search_matches,
        default=default_selected,
        key="search_select"
    )

    for s in selected_search:
        st.session_state["selected_blocks"].add(s)
    for s in default_selected:
        if s not in selected_search and s in search_matches:
            st.session_state["selected_blocks"].discard(s)

    # 类名生成逻辑
    if st.session_state["selected_blocks"]:
        st.markdown("---")
        col1, col2 = st.columns([2, 1])
        with col1:
            prefix = st.text_input("输入类名前缀：", "")
        with col2:
            class_type = st.selectbox("选择类类型：", ["PVO", "RVO", "DTO"])

        class_name = f"{prefix}{class_type}"
        st.markdown(f"#### 🧾 生成类名：`{class_name}`")

        # 构建 Java 类代码
        st.markdown("### 🧾 生成的 Java 类")
        selected_blocks_content = []
        for label in st.session_state["selected_blocks"]:
            file = label.split(" | ")[0]
            match = next(
                (b for b in all_blocks[file] if b.strip().replace('\n', ' ') in label), None
            )
            if match:
                selected_blocks_content.append(match)

        class_code = f"public class {class_name} {{\n" + "\n".join(selected_blocks_content) + "}"
        st.code(class_code, language="java")
